#include "ICM.h"
#include "GaussianProcess.h"
#include "GaussianMRFEnergyDoubleton.h"
#include "EMModelInitialization.h"
#include <cv.h>
#include <cvaux.h>
#include <highgui.h>
#include "cvgabor.h"
#include <stdio.h>
#include <iostream>

using namespace cv;

class ImageParameters {
public:
    typedef Mat_<double> features_vector_t;
    typedef int label_t;
    static int _nbFeature;
    static int _height;
    static int _width;
    static int _numberClass;
    static double _beta;

    inline static int getL(void) {
        return _numberClass;
    }//Number of classes

    inline static int getHeight(void) {
        return _height;
    } //Image size

    inline static int getWidth(void) {
        return _width;
    } //Image size    

    inline static int getN(void) {
        return _nbFeature;
    } //Number of features

    inline static double getBeta(void) {
        return _beta;
    } //Homogeneity
};
int ImageParameters::_nbFeature;
int ImageParameters::_height;
int ImageParameters::_width;
int ImageParameters::_numberClass;
double ImageParameters::_beta;

class ICMParameters {
public:
    static double _stop_criterion;
    
    inline static double getStopCriterion() {
        return _stop_criterion;
    } //Stop criterion
};
double ICMParameters::_stop_criterion;


typedef GaussianProcess<ImageParameters> model_t;
typedef GaussianMRFEnergyDoubleton<ImageParameters> energy_t;
typedef ICM<ImageParameters, energy_t, ICMParameters> optimization_t;

typedef typename ImageParameters::features_vector_t features_vector_t;
typedef typename ImageParameters::label_t label_t;
typedef EModelInitialization<ImageParameters> model_init_t;

int main(int argc, char* argv[]) {
    //Random generator init
    srand((unsigned) time(0));
    //Loading image
    IplImage* img = cvLoadImage(argv[1]);

    int numOrientations = 2;
    IplImage *luv, *gray, *final, *segmented;
    IplImage *gabor[numOrientations];

    if (!img) {
        printf("Could not load %s\n", argv[1]);
        exit(0);
    }

    printf("Width: %d\nHeight: %d\n", img->width, img->height);

    //Init model
    ImageParameters::_nbFeature = 2; //3 = luv feature, 2 = uv features
    ImageParameters::_height = img->height;
    ImageParameters::_width = img->width;
    ImageParameters::_numberClass = atoi(argv[2]); //We work with 5 classes
    ImageParameters::_beta = atof(argv[3]);
    ICMParameters::_stop_criterion = 0.005;

    model_t model;
    optimization_t icm;

    //Set image sizes initialisation
    luv = cvCreateImage(cvGetSize(img), img->depth, img->nChannels);
    printf("depth: %d\n channels: %d\n", img->depth, img->nChannels);
    gray = cvCreateImage(cvGetSize(img), img->depth, 1);
    for (int i = 0; i < numOrientations; ++i)
        gabor[i] = cvCreateImage(cvGetSize(img), img->depth, 1);
    final = cvCreateImage(cvGetSize(img), img->depth, 1);
    
    //LUV
    cvCvtColor(img, luv, CV_RGB2Luv);

    //GRAY
    cvCvtColor(luv, gray, CV_BGR2GRAY);

    //Gabor from GRAY
    /*for (int i = 0; i < numOrientations; i++) {
        CvGabor gaborFilter;
        gaborFilter.Init(PI * (double)i / (double)numOrientations, 1, 2 * PI, sqrt(2.0)/32);
        gaborFilter.conv_img(gray, gabor[i], CV_GABOR_MAG);
    }
    
    for (int x = 0; x < img->width; x++)
        for (int y = 0; y < img->height; y++) {
            int sum = 0;
            for (int i = 0; i < numOrientations; i++) {
                sum += CV_IMAGE_ELEM(gabor[i], uchar, y, x);
                }
            CV_IMAGE_ELEM(final, uchar, y, x) = (uchar)sum / numOrientations;
        }
    cvSaveImage("ICMGabor.jpg", final);*/

    //Create features
    features_vector_t** obs = new features_vector_t*[ImageParameters::getHeight()];
    for (int i = 0; i < ImageParameters::getHeight(); ++i) {
        obs[i] = new features_vector_t[ImageParameters::getWidth()];
        for (int j = 0; j < ImageParameters::getWidth(); ++j) {
            features_vector_t features(ImageParameters::getN(),1);
            //for (int k = 0; k < numOrientations; ++k) {
            //    features(k) = ((double)CV_IMAGE_ELEM(gabor[k], uchar, i, j))  / 255.0;
            //}
            
            //features(0) = ((double) ((uchar*) (luv->imageData + luv->widthStep * i))[j * luv->nChannels]) / 255.0; //L
            features(0) = ((double) ((uchar*) (luv->imageData + luv->widthStep * i))[j * luv->nChannels + 1]) / 255.0; //U
            features(1) = ((double) ((uchar*) (luv->imageData + luv->widthStep * i))[j * luv->nChannels + 2]) / 255.0; //V
            //features(2) = ((double)CV_IMAGE_ELEM(final, uchar, i, j)) / 255.0; //Gabor
            
            obs[i][j] = features;
        }
    }

    //Set gaussian parameters from data (supervised)
    /*Mat_<int> parts(ImageParameters::getL(), 4);
    parts << 0, ImageParameters::getHeight() / 4, 0, ImageParameters::getWidth() / 4,
            0, ImageParameters::getHeight() / 4, (ImageParameters::getWidth()*3) / 4, ImageParameters::getWidth(),
            (ImageParameters::getHeight()*3) / 4, ImageParameters::getHeight(), 0, ImageParameters::getWidth() / 4,
            (ImageParameters::getHeight()*3) / 4, ImageParameters::getHeight(), (ImageParameters::getWidth()*3) / 4, ImageParameters::getWidth(),
            (ImageParameters::getHeight()*3) / 8, (ImageParameters::getHeight()*5) / 8, (ImageParameters::getWidth()*3) / 8, (ImageParameters::getWidth()*5) / 8;

    for (int a = 0; a < ImageParameters::getL(); ++a) {
        features_vector_t mean;
        features_vector_t variance;
        int nsamples = (parts(a, 1)-parts(a, 0))*(parts(a, 3)-parts(a, 2));
        features_vector_t obs_block[nsamples];
        for (int i = parts(a, 0); i < parts(a, 1); ++i) {
            for (int j = parts(a, 2); j < parts(a, 3); j++) {
                obs_block[(i-parts(a, 0))*(parts(a, 3)-parts(a, 2))+(j-parts(a, 2))] = obs[i][j];
            }
        }
        calcCovarMatrix(obs_block, nsamples, variance, mean, CV_COVAR_NORMAL);
        model.setMean(a, mean);
        model.setVariance(a, variance);
    }*/
     
    label_t** labels = new label_t*[ImageParameters::getHeight()];
    for (int i = 0; i < ImageParameters::getHeight(); ++i) {
        labels[i] = new label_t[ImageParameters::getWidth()];
    }

    model_init_t emInit;
    //Init model
    emInit(model, obs, labels);
    for (int a = 0; a < ImageParameters::getL(); ++a) {
        std::cout << "Mean "<< a << " :\n" << model.getMean(a) << std::endl;
        std::cout << "Variance "<< a << " :\n" << model.getVariance(a) << std::endl;
    }

    //Output the initilisation
    segmented = cvCreateImage(cvGetSize(img), img->depth, 1);
    for (int x = 0; x < img->width; x++)
        for (int y = 0; y < img->height; y++) {
            CV_IMAGE_ELEM(segmented, uchar, y, x) = (255 * labels[y][x]) / ImageParameters::getL();
        }
    char buffer [50];
    sprintf(buffer,"ML ICMsegmented n=%d B=%d.jpg",ImageParameters::getL(),(int)ImageParameters::getBeta());
    cvSaveImage(buffer, segmented);
    
    //Segment the image
    icm(model, obs, labels, true);
    //Output the result
    for (int x = 0; x < img->width; x++)
        for (int y = 0; y < img->height; y++) {
            CV_IMAGE_ELEM(segmented, uchar, y, x) = (255 * labels[y][x]) / ImageParameters::getL();
        }
    sprintf(buffer,"ICMsegmented n=%d B=%d.jpg",ImageParameters::getL(),(int)ImageParameters::getBeta());
    cvSaveImage(buffer, segmented);
    
    //Free
    for (int i = 0; i < ImageParameters::getHeight(); ++i) {
        delete[] obs[i];
    }
    delete[] obs;
    for (int i = 0; i < ImageParameters::getHeight(); ++i) {
        delete[] labels[i];
    }
    delete[] labels;

    cvReleaseImage(&img);
    cvReleaseImage(&luv);
    cvReleaseImage(&final);
    cvReleaseImage(&gray);
    cvReleaseImage(&segmented);
    
    for(int i = 0; i < numOrientations; ++i){
        cvReleaseImage(&(gabor[i]));
    }
}
